package com.kongjs.common.sse.core;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class SseEmitterManager {

    private final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    private SseEmitter getSseEmitter(String clientId) {
        SseEmitter sseEmitter = sseEmitterMap.get(clientId);
        if (!ObjectUtils.isEmpty(sseEmitter)) {
            return sseEmitter;
        }
        sseEmitter = new SseEmitter(0L);
        sseEmitter.onTimeout(() -> sseEmitterMap.remove(clientId));
        sseEmitter.onCompletion(() -> sseEmitterMap.remove(clientId));
        sseEmitter.onError((ex) -> sseEmitterMap.remove(clientId));
        sseEmitterMap.put(clientId, sseEmitter);
        return sseEmitter;
    }

    public SseEmitter connect(String clientId) {
        Assert.hasText(clientId, "clientId is null");
        return getSseEmitter(clientId);
    }

    public void disconnect(String clientId) {
        Assert.hasText(clientId, "clientId is null");
        try {
            SseEmitter sseEmitter = sseEmitterMap.get(clientId);
            if (ObjectUtils.isEmpty(sseEmitter)) {
                return;
            }
            sseEmitter.complete();
            sseEmitterMap.remove(clientId);
        } catch (Exception e) {
            log.error("sse disconnect clientId {} error {}", clientId, e.getMessage());
            sseEmitterMap.remove(clientId);
        }
    }

    public void send(String clientId, String eventId, String eventName, String message) {
        Assert.hasText(clientId, "clientId is null");
        Assert.hasText(eventId, "eventId is null");
        Assert.hasText(eventId, "eventName is null");
        Assert.hasText(eventId, "message is null");
        SseEmitter sseEmitter = getSseEmitter(clientId);
        try {
            SseEmitter.SseEventBuilder builder = SseEmitter.event()
                    .id(eventId)
                    .name(eventName)
                    .data(message, MediaType.TEXT_PLAIN)
                    .comment("none")
                    .reconnectTime(3000);
            sseEmitter.send(builder);
        } catch (Exception e) {
            log.error("sse send clientId {} eventId {} eventName {} message {} error {}", clientId, eventId, eventName, message, e.getMessage());
            sseEmitter.completeWithError(e);
        }
    }

    public void sendAll(String eventId, String eventName, String message) {
        Assert.hasText(eventId, "eventId is null");
        Assert.hasText(eventId, "eventName is null");
        Assert.hasText(eventId, "message is null");
        for (Map.Entry<String, SseEmitter> entry : sseEmitterMap.entrySet()) {
            send(entry.getKey(), eventId, eventName, message);
        }
    }
}
